import copy
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F


class PolicyGradient(object):
    def __init__(self, model, lr, device):
        """
        Policy Gradient algorithm
        :param model: (nn.Module) policy的前向网络
        :param lr: (float) 学习率
        """
        self.model = model
        self.model.to(device)
        assert isinstance(lr, float)
        self.lr = lr
        self.device = device
        # 使用Adam优化器
        self.optimizer = optim.Adam(self.model.parameters(), lr=self.lr)

    def predict(self, obs):
        # 使用policy model预测输出的动作概率
        return self.model(obs)

    def learn(self, obs, action, reward):
        # 用policy gradient 算法更新policy model
        act_prob = self.model(obs)  # 获取输出动作概率

        # loss = -R * (log(P) * one_hot(A)) = -R * CrossEntropy(P, A)
        log_prob = torch.sum(-1.0 * torch.log(act_prob) * self.to_one_hot(action, act_prob.shape[1]), dim=1)
        loss = log_prob * reward
        loss = torch.mean(loss)

        self.optimizer.zero_grad()
        loss.backward()
        self.optimizer.step()
        return loss

    def to_one_hot(self, var, dim):
        var_dim = var.shape[0]
        one_hot = torch.zeros((var_dim, dim)).to(self.device)
        one_hot.scatter_(1, var.long().unsqueeze(-1), 1)
        return one_hot
